/*
代码流程：设备检查 → 主机内存注册 → 设备内存分配 → H2D拷贝 → 核函数1（device0） → D2D拷贝 → 核函数2（device1） → D2H拷贝 → H2H拷贝 → 资源清理
*/
#include <iostream>
#include <vector>
#include <cmath>
#include <hip/hip_runtime.h>

// 错误检查宏
#define HIP_CHECK(error)                                                                          \
    if (error != hipSuccess) {                                                                    \
        std::cerr << "HIP error: " << hipGetErrorString(error) << " at line " << __LINE__ << std::endl; \
        exit(EXIT_FAILURE);                                                                       \
    }

// 核函数：对数组中的每个元素进行平方操作
__global__ void squareElements(float* data, size_t n) {
    // 计算三维线程索引
    size_t x = blockIdx.x * blockDim.x + threadIdx.x;
    size_t y = blockIdx.y * blockDim.y + threadIdx.y;
    size_t z = blockIdx.z * blockDim.z + threadIdx.z;

    // 计算线性索引 - 使用标准的三维到一维索引转换
    size_t idx = x + y * blockDim.x * gridDim.x + z * blockDim.x * gridDim.x * blockDim.y * gridDim.y;

    if (idx < n) {
        data[idx] = data[idx] * data[idx];
    }
}

int main() {
    // 检查是否有多个GPU可用
    int deviceCount = 0;
    HIP_CHECK(hipGetDeviceCount(&deviceCount));
    std::cout << "Number of HIP-capable devices: " << deviceCount << std::endl;

    // 如果设备数量小于2，直接退出
    if (deviceCount < 2) {
        std::cerr << "Error: This program requires at least 2 GPUs. Exiting." << std::endl;
        return EXIT_FAILURE;
    }

    // 可调整的参数
    const size_t N = 4096;  // 数据元素数量
    const size_t bytes = N * sizeof(float);

    // 可调整的三维网格和块配置
    dim3 blockSize(8, 4, 2);
    dim3 gridSize(1, 4, 16);

    std::cout << "Launch configuration: Grid(" << gridSize.x << ", " << gridSize.y << ", " << gridSize.z
              << "), Block(" << blockSize.x << ", " << blockSize.y << ", " << blockSize.z
              << "), Data size: " << N << " elements (" << bytes / 1024 << " KB)" << std::endl;

    // 初始化数据
    std::vector<float> h_srcData(N);
    std::vector<float> h_dstData(N);

    for (size_t i = 0; i < N; i++) {
        h_srcData[i] = static_cast<float>(i);
        h_dstData[i] = 0.0f;
    }

    // 1. 使用 hipHostRegister 注册主机内存
    HIP_CHECK(hipHostRegister(h_srcData.data(), bytes, hipHostRegisterDefault));
    HIP_CHECK(hipHostRegister(h_dstData.data(), bytes, hipHostRegisterDefault));

    std::cout << "Host memory successfully registered" << std::endl;

    // 2. 分配设备内存（在不同设备上）
    float *d_srcData, *d_dstData;
    int device0 = 0;
    int device1 = 1 % deviceCount;

    HIP_CHECK(hipSetDevice(device0));
    HIP_CHECK(hipMalloc(&d_srcData, bytes));

    // 在第二个设备上分配内存
    HIP_CHECK(hipSetDevice(device1));
    HIP_CHECK(hipMalloc(&d_dstData, bytes));
    HIP_CHECK(hipSetDevice(device0)); // 切换回原始设备

    std::cout << "Device memory allocated on devices " << device0 << " and " << device1 << " for D2D copy" << std::endl;

    // 3. H2D: 主机到设备的内存拷贝（使用已注册的主机内存）
    HIP_CHECK(hipMemcpy(d_srcData, h_srcData.data(), bytes, hipMemcpyHostToDevice));
    std::cout << "H2D memory copy completed" << std::endl;

    // 4. 在设备上执行计算（使用核函数）
    hipLaunchKernelGGL(squareElements,
                       gridSize, blockSize, 0, 0,
                       d_srcData, N);
    HIP_CHECK(hipDeviceSynchronize());
    std::cout << "Kernel execution completed (squaring operation)" << std::endl;

    // 5. 启用对等访问（如果需要）
    int canAccessPeer = 0;
    HIP_CHECK(hipDeviceCanAccessPeer(&canAccessPeer, device0, device1));

    if (canAccessPeer) {
        HIP_CHECK(hipDeviceEnablePeerAccess(device1, 0));
        std::cout << "Peer access enabled between devices " << device0 << " and " << device1 << std::endl;
    } else {
        std::cerr << "Warning: Peer access not available between devices " << device0 << " and " << device1 << std::endl;
    }

    // 6. D2D: 设备到设备的内存拷贝（在不同设备之间）
    HIP_CHECK(hipMemcpyPeer(d_dstData, device1,
                           d_srcData, device0,
                           bytes));
    std::cout << "D2D memory copy between devices completed" << std::endl;

    // 7. 再次执行核函数（对拷贝后的数据进行操作）
    HIP_CHECK(hipSetDevice(device1)); // 切换到目标设备启动核函数

    hipLaunchKernelGGL(squareElements,
                       gridSize, blockSize, 0, 0,
                       d_dstData, N);
    HIP_CHECK(hipDeviceSynchronize());
    std::cout << "Second kernel execution completed (fourth power operation)" << std::endl;

    // 8. 切换回原始设备
    HIP_CHECK(hipSetDevice(device0));

    // 9. D2H: 设备到主机的内存拷贝（拷贝到已注册的主机内存）
    HIP_CHECK(hipMemcpy(h_dstData.data(), d_dstData, bytes, hipMemcpyDeviceToHost));
    std::cout << "D2H memory copy completed" << std::endl;

    // 10. 在主机上验证结果
    bool correct = true;
    size_t firstErrorIdx = 0;
    float expectedValue = 0.0f, actualValue = 0.0f;

    for (size_t i = 0; i < N; i++) {
        float expected = static_cast<float>(i);
        expected = expected * expected;  // 第一次平方
        expected = expected * expected;  // 第二次平方（四次方）

        if (fabs(h_dstData[i] - expected) > 1e-5) {
            correct = false;
            firstErrorIdx = i;
            expectedValue = expected;
            actualValue = h_dstData[i];
            break;
        }
    }

    if (correct) {
        std::cout << "All operations completed successfully, data verification passed!" << std::endl;
    } else {
        std::cerr << "Data verification failed at index " << firstErrorIdx
                  << ", expected: " << expectedValue << ", got: " << actualValue << std::endl;
    }

    // 11. H2H: 主机到主机的内存拷贝（使用已注册的主机内存）
    // 先重置目标数据
    for (size_t i = 0; i < N; i++) {
        h_dstData[i] = 0.0f;
    }

    HIP_CHECK(hipMemcpy(h_dstData.data(), h_srcData.data(), bytes, hipMemcpyHostToHost));
    std::cout << "H2H memory copy completed" << std::endl;

    // 12. 验证H2H拷贝结果
    correct = true;
    for (size_t i = 0; i < N; i++) {
        if (h_dstData[i] != h_srcData[i]) {
            std::cerr << "H2H copy error at index " << i << ", expected: " << h_srcData[i]
                      << ", got: " << h_dstData[i] << std::endl;
            correct = false;
            break;
        }
    }

    if (correct) {
        std::cout << "H2H memory copy verification passed!" << std::endl;
    } else {
        std::cerr << "H2H memory copy verification failed!" << std::endl;
    }

    // 13. 清理资源
    HIP_CHECK(hipHostUnregister(h_srcData.data()));
    HIP_CHECK(hipHostUnregister(h_dstData.data()));

    // 禁用对等访问（如果已启用）
    if (canAccessPeer) {
        HIP_CHECK(hipDeviceDisablePeerAccess(device1));
        std::cout << "Peer access disabled" << std::endl;
    }

    // 释放设备内存
    HIP_CHECK(hipSetDevice(device0));
    HIP_CHECK(hipFree(d_srcData));

    HIP_CHECK(hipSetDevice(device1));
    HIP_CHECK(hipFree(d_dstData));

    // 切换回原始设备
    HIP_CHECK(hipSetDevice(device0));

    std::cout << "Resources successfully released" << std::endl;

    return 0;
}